package timewheel

/*
 * Time wheel timer
 */

import (
	"container/list"
	"fmt"
	"sync"
	"time"
)

// default slot size.
const gDefaultSlotSize int = 1023

// TimerTask represents a task to be executed
type TimerTask struct {
	// delay time
	Delay time.Duration

	// whether is repeated
	Period bool

	// timer call back
	Callback func()

	// next time
	next time.Time

	// number of rounds before execution.
	rounds int
}

// TimeWheel represents a time wheel which can manager all kinds of timer.
type TimeWheel struct {
	interval   time.Duration
	slots      []*list.List
	ticker     *time.Ticker
	currentPos int
	slotSize   int
	addChan    chan *TimerTask
	removeChan chan *TimerTask
	stopChan   chan struct{}
	taskMap    map[*TimerTask]*list.Element
	mutex      sync.Mutex
}

// NewTimeWheel creates a new TimeWheel
func NewTimeWheel(interval time.Duration, slotSize int) *TimeWheel {
	if interval == 0 {
		interval = time.Second
	}
	if slotSize <= 0 {
		slotSize = gDefaultSlotSize
	}
	tw := &TimeWheel{
		interval:   interval,
		slots:      make([]*list.List, slotSize),
		currentPos: 0,
		slotSize:   slotSize,
		addChan:    make(chan *TimerTask, 1024),
		removeChan: make(chan *TimerTask, 1024),
		stopChan:   make(chan struct{}),
		taskMap:    make(map[*TimerTask]*list.Element),
	}

	for i := 0; i < slotSize; i++ {
		tw.slots[i] = list.New()
	}

	return tw
}

// Start starts the TimeWheel
func (tw *TimeWheel) Start() {
	tw.ticker = time.NewTicker(tw.interval)
	go tw.run()
}

// Stop stops the TimeWheel module.
func (tw *TimeWheel) Stop() {
	close(tw.stopChan)
}

// AddTaskAtTime adds a timer at executeAt time.
func (tw *TimeWheel) AddTaskAtTime(callback func(), executeAt time.Time) (*TimerTask, error) {
	now := time.Now()
	if executeAt.Before(now) {
		return nil, fmt.Errorf("time is before now")
	}

	delay := executeAt.Sub(now)
	task := &TimerTask{
		Delay:    delay,
		Period:   false,
		Callback: callback,
		next:     executeAt,
	}
	tw.AddTask(task)
	return task, nil
}

// AddTask adds a task to the TimeWheel
func (tw *TimeWheel) AddTask(task *TimerTask) {
	if task == nil {
		return
	}
	tw.addChan <- task
}

// AddTaskTimer adds task timer.
func (tw *TimeWheel) AddTaskTimer(delay time.Duration, period bool, callback func()) *TimerTask {
	if delay == 0 {
		delay = time.Second
	}
	task := &TimerTask{
		Delay:    delay,
		Period:   period,
		Callback: callback,
	}
	tw.AddTask(task)
	return task
}

// RemoveTask removes a task from the TimeWheel
func (tw *TimeWheel) RemoveTask(task *TimerTask) {
	if task == nil {
		return
	}
	tw.removeChan <- task
}

func (tw *TimeWheel) run() {
	for {
		select {
		case <-tw.ticker.C:
			tw.tickHandler()
		case task := <-tw.addChan:
			tw.addTask(task)
		case task := <-tw.removeChan:
			tw.removeTask(task)
		case <-tw.stopChan:
			tw.ticker.Stop()
			return
		}
	}
}

func (tw *TimeWheel) tickHandler() {
	tw.mutex.Lock()
	defer tw.mutex.Unlock()

	// Should push to next slot.
	var needDelayTask []*TimerTask = make([]*TimerTask, 0, 64)
	currentSlot := tw.slots[tw.currentPos]
	now := time.Now()
	for e := currentSlot.Front(); e != nil; {
		task := e.Value.(*TimerTask)
		if task.rounds > 0 {
			// rounds还没用到，继续下一个循环。
			task.rounds--
			e = e.Next()
		} else if now.After(task.next) || now.Equal(task.next) {
			// 如果rounds到了，并且时间到了，就执行
			go task.Callback()
			next := e.Next()
			currentSlot.Remove(e)
			delete(tw.taskMap, task)
			e = next

			// 执行完成后
			if task.Period {
				tw.internalAdd(task)
			}
		} else {
			// 时间还没到，放入needDelayTask，等到下次tick.
			needDelayTask = append(needDelayTask, task)

			// 从当前slot中删除。
			next := e.Next()
			currentSlot.Remove(e)
			delete(tw.taskMap, task)
			e = next
		}
	}
	tw.currentPos = (tw.currentPos + 1) % tw.slotSize

	// 把上个tick还没用完成的，放到下一个slot中。
	for _, task := range needDelayTask {
		tw.taskMap[task] = tw.slots[tw.currentPos].PushBack(task)
	}
}

func (tw *TimeWheel) internalAdd(task *TimerTask) {
	if task.Delay < tw.interval {
		task.Delay = tw.interval
	}
	task.next = time.Now().Add(task.Delay)

	virtualID := int(task.Delay / tw.interval)
	pos := (tw.currentPos + virtualID) % tw.slotSize
	task.rounds = virtualID / tw.slotSize

	element := tw.slots[pos].PushBack(task)
	tw.taskMap[task] = element
}

func (tw *TimeWheel) addTask(task *TimerTask) {
	tw.mutex.Lock()
	defer tw.mutex.Unlock()

	tw.internalAdd(task)
}

func (tw *TimeWheel) removeTask(task *TimerTask) {
	tw.mutex.Lock()
	defer tw.mutex.Unlock()

	if element, exists := tw.taskMap[task]; exists {
		for i := 0; i < tw.slotSize; i++ {
			if tw.slots[i].Remove(element) != nil {
				break
			}
		}
		delete(tw.taskMap, task)
	}
}

// FindTask finds a task in the TimeWheel, and find return true, otherwise return false.
func (tw *TimeWheel) FindTask(task *TimerTask) bool {
	tw.mutex.Lock()
	defer tw.mutex.Unlock()

	if _, exists := tw.taskMap[task]; exists {
		return true
	}
	return false
}

// A global time wheel timer object.
var TimeWheelObj *TimeWheel

func init() {
	TimeWheelObj = NewTimeWheel(time.Second, gDefaultSlotSize)
}
